import onnxruntime as ort
import time
from datetime import datetime
from abc import ABC, abstractmethod
from typing import Dict, Optional, List
import numpy as np
import logging

from .yolo_processors import YOLOPreprocessor, YOLOPostprocessor


class ModelWrapper(ABC):
    """模型包装器抽象基类"""
    
    def __init__(self, model_path: str):
        self.model_path = model_path
        self.preprocessor = None
        self.postprocessor = None
        self.session = None  # 延迟初始化
        self._load_model()

    def _load_model(self):
        """
        加载模型的方法。子类可以重写此方法以自定义模型加载逻辑。
        基类实现会尝试从 self.model_path 加载 ONNX 模型。
        """
        if self.model_path and not self.session:
            try:
                self.session = ort.InferenceSession(self.model_path)
            except Exception as e:
                logging.error(f"Failed to load model from {self.model_path}: {e}")
                raise
    
    @abstractmethod
    def predict(self, image: np.ndarray, options: Dict = None) -> Dict:
        """执行预测并返回结果"""
        pass
    
    def generate_report(self, result: Dict, start_time: float) -> Dict:
        """生成报告数据"""
        end_time = time.time()
        report_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        report_duration = round(end_time - start_time, 2)
        report_total_defects = len(result.get('defect_boxes', []))
        report_result = "合格" if report_total_defects == 0 else "不合格"
        
        return {
            "access": not result.get('defect_boxes', []),
            "defect_boxes": result.get('defect_boxes', []),
            "reportTime": report_time,
            "reportDuration": report_duration,
            "reportTotalDefects": report_total_defects,
            "reportResult": report_result,
        }


class YOLOModelWrapper(ModelWrapper):
    """YOLO模型包装器"""
    
    def __init__(
            self, 
            model_path: str,
            preprocessor: Optional[YOLOPreprocessor] = None,
            postprocessor: Optional[YOLOPostprocessor] = None):

        super().__init__(model_path)
        
        # 使用用户提供的预处理器或默认预处理器
        self.preprocessor = preprocessor if preprocessor is not None else YOLOPreprocessor()
        
        # 使用用户提供的后处理器或默认后处理器
        self.postprocessor = postprocessor if postprocessor is not None else YOLOPostprocessor()
    
    def predict(self, image: np.ndarray, options: Dict = None) -> Dict:
        input_data, meta = self.preprocessor.preprocess(image)
        outputs = self.session.run(None, {"images": input_data})
        result = self.postprocessor.postprocess(outputs, meta)
        
        # 添加图像信息到结果中
        result["image"] = {
            "horizontalRatio": meta["ratio_h"],
            "verticalRatio": meta["ratio_w"],
            "originWidth": meta["original_shape"][1],
            "originHeight": meta["original_shape"][0],
            "targetWidth": meta["target_shape"][0],
            "targetHeight": meta["target_shape"][1],
        }
        
        return result


class MutilateModelWrapper(ModelWrapper):
    """多模型包装器，用于组合多个模型的检测结果"""
    
    def __init__(self, 
                 model_configs: List[Dict],
                 preprocessor: Optional[YOLOPreprocessor] = None,
                 postprocessor: Optional[YOLOPostprocessor] = None,
                 stop_on_defect: bool = False):
        """
        初始化多模型包装器
        
        Args:
            model_configs: 模型配置列表，每个配置包含模型类型、路径和名称
            preprocessor: 预处理器实例
            postprocessor: 后处理器实例
            stop_on_defect: 检测到缺陷时是否中断后续检测
        """
        # 由于继承自 ModelWrapper，需要传入一个模型路径，但这里我们不使用它
        super().__init__(model_path=None)
        
        self.model_configs = model_configs
        self.models = []
        self.stop_on_defect = stop_on_defect
        
        # 使用用户提供的预处理器或默认预处理器
        self.preprocessor = preprocessor if preprocessor is not None else YOLOPreprocessor()
        
        # 使用用户提供的后处理器或默认后处理器
        self.postprocessor = postprocessor if postprocessor is not None else YOLOPostprocessor()
    
    def addModels(self, name, model, defect_types):
        """初始化所有配置的模型"""
        self.models.append({
            "name": name,
            "model": model,
            "defect_types": defect_types
        })
    
    def predict(self, image: np.ndarray, options: Dict = None) -> Dict:
        """执行多模型预测并合并结果"""
        all_results = {"defect_boxes": [], "image": {}, "models": []}
        
        # 执行预处理器（只需要执行一次，所有模型使用相同的预处理结果）
        input_data, meta = self.preprocessor.preprocess(image)
        
        has_defect = False
        
        for model_info in self.models:
            model_name = model_info["name"]
            model = model_info["model"]
            defect_types = model_info["defect_types"]
            
            # 使用相同的预处理结果
            model.preprocessor = self.preprocessor
            model.postprocessor = self.postprocessor
            
            # 执行预测
            start_time = time.time()
            result = model.predict(input_data)
            end_time = time.time()
            
            # 检查是否有缺陷
            current_has_defect = len(result.get("defect_boxes", [])) > 0
            
            # 记录模型结果
            model_result = {
                "model_name": model_name,
                "defect_types": defect_types,
                "has_defect": current_has_defect,
                "defect_boxes": result.get("defect_boxes", []),
                "processing_time": end_time - start_time
            }
            
            all_results["models"].append(model_result)
            
            # 合并缺陷框
            if current_has_defect:
                # 为每个缺陷框添加来源模型信息
                for box in model_result["defect_boxes"]:
                    box["source_model"] = model_name
                    box["defect_types"] = defect_types
                    
                all_results["defect_boxes"].extend(model_result["defect_boxes"])
                has_defect = True
            
            # 如果设置了检测到缺陷就中断，并且当前模型检测到了缺陷，则停止后续检测
            if self.stop_on_defect and current_has_defect:
                break
        
        # 添加图像信息
        all_results["image"] = {
            "horizontalRatio": meta["ratio_h"],
            "verticalRatio": meta["ratio_w"],
            "originWidth": meta["original_shape"][1],
            "originHeight": meta["original_shape"][0],
            "targetWidth": meta["target_shape"][0],
            "targetHeight": meta["target_shape"][1],
        }
        
        # 添加总体检测结果
        all_results["has_defect"] = has_defect
        all_results["total_defects"] = len(all_results["defect_boxes"])
        
        return all_results
    
    def generate_report(self, result: Dict, start_time: float) -> Dict:
        """生成多模型检测报告"""
        end_time = time.time()
        report_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        report_duration = round(end_time - start_time, 2)
        
        # 计算每个模型的检测结果
        model_reports = []
        total_defects = 0
        
        for model_result in result.get("models", []):
            model_report = {
                "model_name": model_result["model_name"],
                "defect_types": model_result["defect_types"],
                "has_defect": model_result["has_defect"],
                "defect_count": len(model_result["defect_boxes"]),
                "processing_time": round(model_result["processing_time"], 4)
            }
            model_reports.append(model_report)
            total_defects += model_report["defect_count"]
        
        report_result = "合格" if total_defects == 0 else "不合格"
        
        return {
            "access": total_defects == 0,
            "defect_boxes": result.get("defect_boxes", []),
            "reportTime": report_time,
            "reportDuration": report_duration,
            "reportTotalDefects": total_defects,
            "reportResult": report_result,
            "modelReports": model_reports
        }

